import pandas as pd
from Bio import SeqIO
from Bio import pairwise2
from Bio.Align import substitution_matrices

def smith_waterman(seq1, seq2, gap_open_penalty=-10, gap_extension_penalty=-0.5):
    matrix = substitution_matrices.load("BLOSUM62")
    alignments = pairwise2.align.localds(seq1, seq2, matrix, gap_open_penalty, gap_extension_penalty)
    best_score = alignments[0][2]
    short_seq = min(seq1, seq2, key=len)
    max_possible_score = sum(matrix.get((res, res), 0) for res in short_seq)
    return (best_score / max_possible_score) * 100

def drop_results(df_path, models):
    df = pd.read_csv(df_path, sep="\t")
    try:
        df = df.drop(models, axis=1)
    except:
        return 
    df.to_csv(df_path, sep="\t")


def read_fasta(file_path, key):
    return str(getattr(SeqIO.read(file_path, 'fasta'), key))


def mutant_filter(df, mutant_site=0):
    rows = ''.join(df["mutant"].tolist())
    sep = ":"
    if ";" in rows:
        sep = ";"
    if mutant_site:
        trg_df = df.loc[df.apply(lambda row: len(row["mutant"].split(sep))==mutant_site, axis=1)]
    else:
        trg_df = df
    return trg_df


def scan_max_mutant(df, seq, max_len=1022):
    # df: dataframe of the mutant file
    df_single = mutant_filter(df, 1)
    mutants = list(df_single["mutant"])
    
    # WT in the first row
    has_wt = False
    if mutants[0].lower() == "wt":
        mutants = mutants[1:]
        df_no_wt = df_single[1:]
        has_wt = True
    
    mutant_pos = [int(m[1:-1]) - 1 for m in mutants]
    max_pos, min_pos = max(mutant_pos), min(mutant_pos)
    
    if max_pos < max_len:
        return df, seq[:max_len], 1
     
    if max_pos - min_pos + 1 <= max_len:
        seq_left_to_right = seq[min_pos: min_pos + max_len]
        # contain the last mutation residue
        seq_right_to_left = seq[max_pos - max_len + 1: max_pos + 1]
        
        # select the longer sequence
        if len(seq_left_to_right) > len(seq_right_to_left):
            seq = seq_left_to_right
            offset = min_pos + 1
            if has_wt:
                df_bool = df_no_wt.apply(lambda row: int(row["mutant"][1:-1]) >= offset and int(row["mutant"][1:-1]) < (min_pos + max_len + 1), axis=1)
                df_bool.loc[0] = True
                df = df.loc[df_bool]
            else:
                df = df.loc[df.apply(lambda row: int(row["mutant"][1:-1]) >= offset and int(row["mutant"][1:-1]) < (min_pos + max_len + 1), axis=1)]
        else:
            seq = seq_right_to_left
            offset = max_pos - max_len + 2
            if has_wt:
                df_bool = df_no_wt.apply(lambda row: int(row["mutant"][1:-1]) >= offset and int(row["mutant"][1:-1]) < max_pos + 2, axis=1)
                df_bool.loc[0] = True
                df = df.loc[df_bool]
            else:
                df = df.loc[df.apply(lambda row: int(row["mutant"][1:-1]) >= offset and int(row["mutant"][1:-1]) < max_pos + 2, axis=1)]
        return df, seq, offset
    
    global_count = 0
    global_left, global_right = min_pos, max_len
    window_left, window_right = min_pos, max_len
    # scan the whole sequence
    while(window_right < len(seq)):
        window_count = len([pos for pos in mutant_pos if pos >= window_left and pos < window_right])
        if window_count > global_count:
            global_left, global_right = window_left, window_right
            global_count = window_count
        window_left += 1
        window_right += 1
    
    if has_wt:
        df_bool = df_no_wt.apply(lambda row: int(row["mutant"][1:-1]) >= global_left+1 and int(row["mutant"][1:-1]) <= global_right, axis=1)
        df_bool.loc[0] = True
        df = df.loc[df_bool]
    else:
        df = df.loc[df.apply(lambda row: int(row["mutant"][1:-1]) >= global_left+1 and int(row["mutant"][1:-1]) <= global_right, axis=1)]
    
    # if the sequence is too short, extend it to the max_len
    if global_right - global_left + 1 < max_len:
        global_left = global_right - max_len
     
    return df, seq[global_left:global_right], global_left +1    


def full_sequence(origin_sequence, raw_mutant, offset):
    # if the mutant is in the format of "A123V;A124V"
    sep = ";"
    # if the mutant is in the format of "A123V:A124V"
    if ":" in raw_mutant:
        sep = ":"
        
    sequence = origin_sequence
    for raw_mut in raw_mutant.split(sep):
        # if the mutant is "WT"
        if raw_mut.lower() == "wt":
            continue
        to = raw_mut[-1]
        pos = int(raw_mut[1:-1]) - offset
        assert sequence[pos] == raw_mut[0], print(sequence[pos], raw_mut[0], pos)
        sequence = sequence[:pos] + to + sequence[pos + 1:]
    return sequence